# !/usr/bin/env python
# -*- coding:utf-8 -*-
# @Project   :pyTools 
# @FileName  :socket_ex.py
# @Time      :2023/12/1 13:10
# @Author    :zxc19
# @IDE       :PyCharm
# @DESC      :
import inspect
import json
import logging

from locust import task, constant, between, constant_throughput, constant_pacing, env, User, tag, events
from locust.runners import MasterRunner, WorkerRunner, LocalRunner
from locust.argument_parser import LocustArgumentParser
from locust.exception import LocustError
from users.socket import UDPSocketAbstractUser
from utils.ip_utils import get_local_ips


def show(user: User):
    user.environment.events.request.fire(
        request_type="外部任务",
        name='show',
        response_time=None,
        response_length=0,
        exception=None,
        context={},
    )


def show2(user: User):
    user.environment.events.request.fire(
        request_type="外部任务",
        name='show2',
        response_time=None,
        response_length=0,
        exception=None,
        context={},
    )


@events.init_command_line_parser.add_listener
def _(parser: LocustArgumentParser):
    """
    添加locust命令行参数
    :param parser:
    :return:
    """

    config_info = [
        {
            "send_host": "10.121.110.249",
            "dst_ip": "10.121.110.42",
            "dst_port": 514
        },
        {
            "send_host": "2003::110:220",
            "dst_ip": "2003::110:42",
            "dst_port": 514
        },
        {
            "send_host": "10.121.111.177",
            "dst_ip": "10.121.110.161",
            "dst_port": 514
        }
    ]
    parser.add_argument(
        '--config-info',  # 参数名称
        env_var='LOCUST_CONFIG_INFO',  # 在环境变量中的名称
        default=json.dumps(config_info),  # 默认值
        include_in_web_ui=True,  # 是否包含在WEB UI内
        help='配置信息JSON'  # 帮助信息
    )
    parser.add_argument(
        '--dst-info',  # 参数名称
        env_var='LOCUST_DST_INFO',  # 在环境变量中的名称
        default=None,  # 默认值
        include_in_web_ui=False,  # 是否包含在WEB UI内
        help='目的信息'  # 帮助信息
    )
    parser.add_argument(
        "--data-type",
        choices=["str", "file"],
        default="file",
        help="数据类型（字符串还是文件），文件时为读取data目录下的文件"
    )
    parser.add_argument(
        "--data",
        type=str,
        env_var="LOCUST_DATA",
        default="test.log",
        help="测试日志数据"
    )


@events.init.add_listener
def on_locust_init(environment: env.Environment, runner: WorkerRunner, **kwargs):
    """
    每个locust进程开始时触发,主要用于分布式负载测试
    :param environment:
    :param runner:
    :param kwargs:
    :return:
    """
    if isinstance(environment.runner, MasterRunner):
        logging.info("初始化Locust负载测试主节点")
    elif isinstance(environment.runner, WorkerRunner):
        logging.info("初始化Locust负载测试从节点")
    elif isinstance(environment.runner, LocalRunner):
        logging.info("初始化Locust单节点负载测试")
    logging.debug(f'环境变量信息==>{environment.parsed_options}')
    logging.debug(f'运行者信息==>{runner}')
    logging.debug(f'其他信息==>{kwargs}')






@events.test_start.add_listener
def on_test_start(environment: env.Environment, **kwargs):
    """
    负载测试开始时执行
    :param environment:
    :param kwargs:
    :return:
    """
    logging.info("开始Locust负载测试")
    logging.debug(f'环境变量信息==>{environment.parsed_options}')
    logging.debug(f'其他信息==>{kwargs}')
    if isinstance(environment.runner, WorkerRunner) or isinstance(environment.runner, LocalRunner):
        local_ips = get_local_ips()
        for config_info in json.loads(environment.parsed_options.config_info):
            if config_info.get('send_host') in local_ips:
                environment.parsed_options.dst_info = config_info
                break
        # 根据数据类型决定数据使用字符串还是文件读取
        if environment.parsed_options.data_type == 'file':
            try:
                with open(f'data/{environment.parsed_options.data}', encoding='utf-8') as f:
                    environment.parsed_options.data = f.read()
            except FileNotFoundError as e:
                logging.error(f'{environment.parsed_options.data}文件不存在,结束测试')
                environment.runner.quit()

def on_test_stop(environment: env.Environment, **kwargs):
    """
    负载测试停止时执行
    :param environment:
    :param kwargs:
    :return:
    """
    logging.info("停止Locust负载测试")
    logging.debug(f'环境变量信息==>{environment}')
    logging.debug(f'其他信息==>{kwargs}')


@events.request.add_listener
def my_request_handler(
        **kwargs
):
    """
    自定义事件监听(参数是environment.events.request.fire里注册的信息)
    :param request_type: 请求类型
    :param name: 事件名称
    :param response_time: 响应事件
    :param response_length: 响应长度
    :param context: 上下文
    :param exception: 异常
    :param kwargs:
    :return:
    """
    # if exception:
    #     logging.error(f"事件请求[{name}]异常==>{exception}")
    # else:
    #     logging.debug(f"事件请求[{name}]成功")
    logging.debug(f'请求捕捉')
    logging.debug(f'其他信息==>{kwargs}')


class UDPSyslogUser(UDPSocketAbstractUser):
    # wait_time = constant(2)  # 每次运行后停顿固定时间(s)
    # wait_time = between(1, 5)  # 每次运行后停顿 在最小值和最大值之间选取一个随机时间(s)
    wait_time = constant_throughput(1)  # 任务每秒运行(最多)X次的自适应时间
    # wait_time = constant_pacing(2)  # 任务每X秒运行(最多)一次的自适应时间

    # last_wait_time = 0
    #
    # def wait_time(self):
    #     """
    #     自定义等待时间
    #     :return:
    #     """
    #     self.last_wait_time += 1
    #     return self.last_wait_time

    # weight = 1  # 用户权重,当存在多个用户类时才会有用
    # fixed_count = 1  # 用户数量,weight属性将被忽略，并将生成确切的用户数.同时命令或web ui上设置的用户数也会失效

    # 任务属性(外部任务列表或外部任务字典,值为权重),用户指定外部的任务,这些任务的首个参数必须是用户实例
    # tasks = [show]
    # tasks = {show: 1, show2: 2}

    # 环境属性
    environment: env.Environment

    def __init__(self, environment):
        super().__init__(environment)
        logging.debug(f'初始化用户:{self.__class__.__name__}')
        # 从环境变量提取配置数据
        self.config_info = self.environment.parsed_options.dst_info
        self.data_type = self.environment.parsed_options.data
        self.data = self.environment.parsed_options.data

    # 标记为任务
    @task
    def syslog(self):
        logging.debug(f'任务:{inspect.stack()[0].function}')
        self.send(self.config_info.get('dst_ip'), self.config_info.get('dst_port'), self.data)


if __name__ == "__main__":
    run_code = 0
